{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test DualKuka in Comp vgrp Env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import time, sys\n",
    "sys.path.append('..')\n",
    "from pb_diff_envs.environment.static.comp_kuka_env import ComposedDynKukaEnv\n",
    "\n",
    "import numpy as np\n",
    "from importlib import reload\n",
    "import matplotlib.pyplot as plt\n",
    "import pybullet as p\n",
    "import pb_diff_envs.objects.static.voxel_group as vv\n",
    "reload(vv)\n",
    "from pb_diff_envs.environment.static.kuka_rand_vgrp import Kuka_VoxelRandGroupList\n",
    "\n",
    "mazelist_config = dict(planner_timeout=30, \n",
    "                        interp_density=1.5,\n",
    "                       min_episode_distance=4*np.pi)\n",
    "\n",
    "import gym\n",
    "env_name = 'Compdualkuka14d-testOnly-rand-v12' \n",
    "# env_name = 'Compdualkuka14d-rand-nv5nv3-se1005-vr0702-hExt22-v0'\n",
    "dvg_group = gym.make(env_name, load_mazeEnv=False, gen_data=True, is_eval=True) # True\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## simple vis of env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "if p.isConnected():\n",
    "    p.disconnect()\n",
    "env = dvg_group.create_single_env(3) # \n",
    "env.unload_env()\n",
    "env.load(GUI=False) # important not delete\n",
    "p.getContactPoints()\n",
    "p.stepSimulation()\n",
    "plt.imshow(env.render())\n",
    "plt.show()\n",
    "\n",
    "p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0)  # Disable the default GUI controls\n",
    "p.configureDebugVisualizer(p.COV_ENABLE_MOUSE_PICKING, 1)  # Enable mouse picking for camera control\n",
    "p.addUserDebugLine([0, 0, 0], [1, 0, 0], lineColorRGB=[1, 0, 0], lineWidth=6)\n",
    "p.addUserDebugLine([0, 0, 0], [0, 1, 0], lineColorRGB=[0, 1, 0], lineWidth=6)\n",
    "p.addUserDebugLine([0, 0, 0], [0, 0, 1], lineColorRGB=[0, 0, 1], lineWidth=6)\n",
    "no_while = True ## can be True when with\n",
    "while True:\n",
    "    p.stepSimulation()\n",
    "    # p.performCollisionDetection()\n",
    "    c_pts = p.getContactPoints()\n",
    "    # print(p.get)\n",
    "    time.sleep(0.1)\n",
    "    if c_pts is not None and len(c_pts) > 0:\n",
    "        # print(p.getContactPoints())\n",
    "        pass\n",
    "        # time.sleep(2)\n",
    "        # p.resetJointStatesMultiDof(env.robot_id, np.arange(7), [[0.]]*7)\n",
    "    if no_while:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.robot.print_joint_states()\n",
    "## print(env.robot.item_id)\n",
    "print('robot_id:', env.robot_id) # list\n",
    "print('limits_high:', env.robot.limits_high)\n",
    "# env.robot.limits_low\n",
    "print(env._get_joint_pose())\n",
    "print(f'{type(env._get_joint_pose())}')\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample a problem/episode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from pb_diff_envs.utils.utils import DotDict\n",
    "from tqdm import tqdm\n",
    "env.load(GUI=False) # False\n",
    "prev_pose = np.array([0.0] * env.robot.config_dim)\n",
    "env.robot.set_config(prev_pose)\n",
    "solutions = []\n",
    "n_traj = 5 # 20\n",
    "for i in tqdm(range(n_traj)):\n",
    "    solution, _ = env.sample_1_episode(prev_pose)\n",
    "    solution = np.array(solution)\n",
    "    prev_pose = solution[-1]\n",
    "    solutions.append(solution)\n",
    "\n",
    "# solutions = np.concatenate(solutions, axis=0)\n",
    "# result = DotDict(solution=solutions)\n",
    "# print(solutions.shape)\n",
    "# print(len(result.solution))\n",
    "# result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import pybullet as p\n",
    "import pybullet_data\n",
    "import numpy as np\n",
    "from pb_diff_envs.utils.utils import save_gif\n",
    "from IPython.display import HTML\n",
    "import base64\n",
    "from pb_diff_envs.objects.trajectory import WaypointLinearTrajectory, WaypointProportionTrajectory\n",
    "from pb_diff_envs.utils.kuka_utils_luo import SolutionInterp\n",
    "# Visualization\n",
    "from importlib import reload\n",
    "import pb_diff_envs.utils.robogroup_utils_luo  as rul; reload(rul)\n",
    "from pb_diff_envs.utils.robogroup_utils_luo import robogroup_visualize_traj_luo\n",
    "from pb_diff_envs.utils.kuka_utils_luo import visualize_kuka_traj_luo\n",
    "\n",
    "\n",
    "# pick one problem\n",
    "sol_interp = SolutionInterp(density=1.5) # mazelist_config['interp_density']\n",
    "sol_i = solutions[0] # which motion trajectory to visualize\n",
    "result = DotDict(solution=sol_i)\n",
    "print(result.solution.shape)\n",
    "traj_vis = sol_interp(result.solution,)\n",
    "\n",
    "gifs, ds, vis_dict = robogroup_visualize_traj_luo(env, traj_vis, is_debug=False)\n",
    "\n",
    "\n",
    "gif_name = './test_comp_dualkuka_wgrp.gif'\n",
    "save_gif(gifs, gif_name, duration=ds)\n",
    "b64 = base64.b64encode(open(gif_name, 'rb').read()).decode('ascii')\n",
    "display(HTML(f'<img src=\"data:image/gif;base64,{b64}\" />'))\n",
    "\n",
    "print(traj_vis.shape)\n",
    "print(result.solution.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.wall_locations\n",
    "# while True:\n",
    "#     time.sleep(0.1)\n",
    "#     p.stepSimulation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  },
  "vscode": {
   "interpreter": {
    "hash": "c1daed7e4d34e14fbb9bb260207f2e4074aea67593e999acdc484a19561ddc79"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
